package com.maverick.util;

import com.google.gson.JsonParser;
import org.apache.flink.core.io.SimpleVersionedSerializer;
import org.apache.flink.streaming.api.functions.sink.filesystem.BucketAssigner;
import org.apache.flink.streaming.api.functions.sink.filesystem.bucketassigners.SimpleVersionedStringSerializer;
import org.apache.flink.util.Preconditions;

import java.time.Instant;
import java.time.ZoneId;
import java.time.format.DateTimeFormatter;

public class JsonTimeFieldBucketAssigner implements BucketAssigner<java.lang.String, String> {

    private static final long serialVersionUID = 1L;

    private static final String DEFAULT_FORMAT_STRING = "yyyy/MM/dd/HH";

    private final String formatString;

    private final ZoneId zoneId;

    private transient DateTimeFormatter dateTimeFormatter;

    private final String timeField;

    public JsonTimeFieldBucketAssigner(String timeField, String formatString) {
        this(timeField, formatString, ZoneId.systemDefault());
    }

    public JsonTimeFieldBucketAssigner(String timeField, String formatString, ZoneId zoneId) {
        this.timeField = timeField;
        this.formatString = Preconditions.checkNotNull(formatString);
        this.zoneId = Preconditions.checkNotNull(zoneId);
    }

    @Override
    public String getBucketId(java.lang.String element, Context context) {
        if (dateTimeFormatter == null) {
            dateTimeFormatter = DateTimeFormatter.ofPattern(formatString).withZone(zoneId);
        }
        long time = timeField.isEmpty() ?
                context.currentProcessingTime()
                : new JsonParser().parse(element).getAsJsonObject().get(timeField).getAsLong();
        return dateTimeFormatter.format(Instant.ofEpochMilli(time));
    }

    @Override
    public SimpleVersionedSerializer<String> getSerializer() {
        return SimpleVersionedStringSerializer.INSTANCE;
    }
}